package com.zzedu.framework.ratelimit;


import com.google.common.collect.Maps;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

@Component
public class RedisSlidingWindowRateLimiter implements ApplicationRunner {

    @Autowired
    RedisTemplate<String, Object> redisTemplate;

    DefaultRedisScript<Long> redisScript;




    // Lua script for sliding window rate limiting
    private static final String LUA_SCRIPT =
            "local key = KEYS[1]\n" +
                    "local limit = tonumber(ARGV[1])\n" +
                    "local window = tonumber(ARGV[2])\n" +
                    "local now = tonumber(ARGV[3])\n" +
                    "local clearBefore = now - window\n" +
                    "redis.call('ZREMRANGEBYSCORE', key, '-inf', clearBefore)\n" +
                    "local currentCount = redis.call('ZCARD', key)\n" +
                    "if currentCount < limit then\n" +
                    "  redis.call('ZADD', key, now, now)\n" +
                    "  redis.call('EXPIRE', key, window)\n" +
                    "  return 1\n" +
                    "else\n" +
                    "  return 0\n" +
                    "end";

    // 存储各个 API 的限流配置
    private Map<String, RateLimitConfig> rateLimitConfigMap = Maps.newConcurrentMap();

    @Override
    public void run(ApplicationArguments args) throws Exception {
        init();
        redisScript = new DefaultRedisScript<Long>(LUA_SCRIPT, Long.class);
    }

    public static class RateLimitConfig {
        private int limit;
        private Duration duration;

        public RateLimitConfig(int limit, Duration duration) {
            this.limit = limit;
            this.duration = duration;
        }

        public int getLimit() {
            return limit;
        }

        public Duration getDuration() {
            return duration;
        }

        public void setLimit(int limit) {
            this.limit = limit;
        }

        public void setDuration(Duration duration) {
            this.duration = duration;
        }
    }

    public void init() {
        rateLimitConfigMap.put("default", new RateLimitConfig(1000, Duration.ofSeconds(1)));
    }

    public boolean isAllowed(String key) {
        RateLimitConfig config = rateLimitConfigMap.getOrDefault(key, rateLimitConfigMap.get("default"));
        int limit = config.getLimit();
        long duration = config.getDuration().toMillis();
        long now = System.currentTimeMillis();
        redisScript.setScriptText(LUA_SCRIPT);
        redisScript.setResultType(Long.class);
        Long result = redisTemplate.execute(redisScript, Collections.singletonList(key), limit, duration, now);
        return result != null && result == 1;
    }

    public void updateRateLimitConfig(String key, int limit, Duration duration) {
        rateLimitConfigMap.put(key, new RateLimitConfig(limit, duration));
    }

    public void updateRateLimitConfig(String key, RateLimitConfig rateLimitConfig) {
        rateLimitConfigMap.put(key, rateLimitConfig);
    }
}
